package ublk

import (
	"context"
	"fmt"
	"os"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/containerd/containerd/log"
	"onyx-snapshotter/pkg/config"
	"onyx-snapshotter/pkg/metrics"
)

// Driver manages ublk block devices
type Driver struct {
	config  *config.UblkConfig
	devices map[int]*Device
	mu      sync.RWMutex
	nextID  int32
}

type Device struct {
	id         int
	devPath    string
	size       int64
	blockSize  int64
	dataSource DataSource

	queues  []*Queue
	running atomic.Bool
	stopCh  chan struct{}
}

type Queue struct {
	id     int
	depth  int
	ring   *IoUring
	stopCh chan struct{}
}

// DataSource interface for providing block data
type DataSource interface {
	ReadAt(ctx context.Context, buf []byte, offset int64) (int, error)
	WriteAt(ctx context.Context, buf []byte, offset int64) (int, error)
	Size() int64
}

// IoUring represents io_uring for async I/O
type IoUring struct {
	fd        int
	sqEntries uint32
	cqEntries uint32
	// Simplified - real implementation would have ring buffers
}

const (
	UBLK_CMD_ADD_DEV   = 0x4004
	UBLK_CMD_DEL_DEV   = 0x4005
	UBLK_CMD_START_DEV = 0x4006
	UBLK_CMD_STOP_DEV  = 0x4007

	UBLK_IO_OP_READ    = 0
	UBLK_IO_OP_WRITE   = 1
	UBLK_IO_OP_FLUSH   = 2
	UBLK_IO_OP_DISCARD = 3
)

func NewDriver(cfg *config.UblkConfig) (*Driver, error) {
	if !cfg.Enabled {
		return nil, fmt.Errorf("ublk is disabled")
	}

	// Check if ublk is available (Linux 6.0+)
	if !isUblkAvailable() {
		if cfg.FallbackToFuse {
			log.L.Warn("ublk not available, would fallback to FUSE")
			return nil, fmt.Errorf("ublk not available")
		}
		return nil, fmt.Errorf("ublk not available and fallback disabled")
	}

	d := &Driver{
		config:  cfg,
		devices: make(map[int]*Device),
		nextID:  0,
	}

	log.L.WithField("num_queues", cfg.NumQueues).Info("ublk driver initialized")
	return d, nil
}

// CreateDevice creates a new ublk block device
func (d *Driver) CreateDevice(ctx context.Context, size int64, source DataSource) (*Device, error) {
	d.mu.Lock()
	defer d.mu.Unlock()

	id := int(atomic.AddInt32(&d.nextID, 1))

	dev := &Device{
		id:         id,
		devPath:    fmt.Sprintf("/dev/ublkb%d", id),
		size:       size,
		blockSize:  4096,
		dataSource: source,
		queues:     make([]*Queue, d.config.NumQueues),
		stopCh:     make(chan struct{}),
	}

	// Create ublk control device
	if err := dev.create(ctx, d.config); err != nil {
		return nil, fmt.Errorf("failed to create device: %w", err)
	}

	// Initialize queues
	for i := 0; i < d.config.NumQueues; i++ {
		queue, err := newQueue(i, d.config.QueueDepth)
		if err != nil {
			dev.destroy()
			return nil, fmt.Errorf("failed to create queue %d: %w", i, err)
		}
		dev.queues[i] = queue
	}

	d.devices[id] = dev
	metrics.UblkDevices.Inc()

	log.G(ctx).WithFields(map[string]interface{}{
		"id":       id,
		"dev_path": dev.devPath,
		"size":     size,
	}).Info("created ublk device")

	return dev, nil
}

// Start starts the ublk device
func (d *Device) Start(ctx context.Context) error {
	if !d.running.CompareAndSwap(false, true) {
		return fmt.Errorf("device already running")
	}

	// Start queue workers
	for _, queue := range d.queues {
		go d.queueWorker(ctx, queue)
	}

	// Notify kernel to start device
	if err := d.startDevice(); err != nil {
		d.running.Store(false)
		return err
	}

	log.G(ctx).WithField("dev_path", d.devPath).Info("started ublk device")
	return nil
}

// Stop stops the ublk device
func (d *Device) Stop(ctx context.Context) error {
	if !d.running.CompareAndSwap(true, false) {
		return fmt.Errorf("device not running")
	}

	close(d.stopCh)

	// Stop all queues
	for _, queue := range d.queues {
		close(queue.stopCh)
	}

	// Notify kernel to stop device
	if err := d.stopDevice(); err != nil {
		log.G(ctx).WithError(err).Warn("failed to stop device")
	}

	log.G(ctx).WithField("dev_path", d.devPath).Info("stopped ublk device")
	return nil
}

// Remove removes the ublk device
func (d *Driver) RemoveDevice(ctx context.Context, id int) error {
	d.mu.Lock()
	defer d.mu.Unlock()

	dev, ok := d.devices[id]
	if !ok {
		return fmt.Errorf("device %d not found", id)
	}

	if dev.running.Load() {
		dev.Stop(ctx)
	}

	if err := dev.destroy(); err != nil {
		return err
	}

	delete(d.devices, id)
	metrics.UblkDevices.Dec()

	log.G(ctx).WithField("id", id).Info("removed ublk device")
	return nil
}

// Queue worker processes I/O requests
func (d *Device) queueWorker(ctx context.Context, queue *Queue) {
	log.G(ctx).WithField("queue_id", queue.id).Debug("queue worker started")

	for {
		select {
		case <-queue.stopCh:
			return
		default:
			// Process I/O requests from io_uring
			if err := d.processRequests(ctx, queue); err != nil {
				log.G(ctx).WithError(err).Error("failed to process requests")
				time.Sleep(time.Millisecond * 10)
			}
		}
	}
}

func (d *Device) processRequests(ctx context.Context, queue *Queue) error {
	// Simplified implementation - real code would use io_uring
	// to receive requests from kernel

	// Mock: simulate receiving a read request
	req := &IoRequest{
		Op:     UBLK_IO_OP_READ,
		Offset: 0,
		Length: 4096,
		Buffer: make([]byte, 4096),
	}

	start := time.Now()
	var err error

	switch req.Op {
	case UBLK_IO_OP_READ:
		_, err = d.dataSource.ReadAt(ctx, req.Buffer, req.Offset)
		metrics.RecordUblkRequest("read", err == nil)
	case UBLK_IO_OP_WRITE:
		_, err = d.dataSource.WriteAt(ctx, req.Buffer, req.Offset)
		metrics.RecordUblkRequest("write", err == nil)
	case UBLK_IO_OP_FLUSH:
		// Flush operation
		metrics.RecordUblkRequest("flush", true)
	}

	duration := time.Since(start).Seconds()
	metrics.UblkRequestDuration.WithLabelValues(opTypeString(req.Op)).Observe(duration)

	return err
}

// Device lifecycle methods

func (d *Device) create(ctx context.Context, cfg *config.UblkConfig) error {
	// Open ublk control device
	fd, err := syscall.Open("/dev/ublk-control", syscall.O_RDWR, 0)
	if err != nil {
		return fmt.Errorf("failed to open ublk control: %w", err)
	}
	defer syscall.Close(fd)

	// IOCTL to create device
	// Simplified - real implementation would use proper ioctl
	return nil
}

func (d *Device) startDevice() error {
	// IOCTL to start device
	return nil
}

func (d *Device) stopDevice() error {
	// IOCTL to stop device
	return nil
}

func (d *Device) destroy() error {
	// IOCTL to delete device
	return nil
}

// Queue initialization

func newQueue(id, depth int) (*Queue, error) {
	ring, err := setupIoUring(uint32(depth))
	if err != nil {
		return nil, err
	}

	return &Queue{
		id:     id,
		depth:  depth,
		ring:   ring,
		stopCh: make(chan struct{}),
	}, nil
}

func setupIoUring(entries uint32) (*IoUring, error) {
	// Simplified - real implementation would setup io_uring
	return &IoUring{
		sqEntries: entries,
		cqEntries: entries * 2,
	}, nil
}

// Helper types and functions

type IoRequest struct {
	Op     int
	Offset int64
	Length int64
	Buffer []byte
}

func isUblkAvailable() bool {
	// Check if /dev/ublk-control exists
	_, err := os.Stat("/dev/ublk-control")
	return err == nil
}

func opTypeString(op int) string {
	switch op {
	case UBLK_IO_OP_READ:
		return "read"
	case UBLK_IO_OP_WRITE:
		return "write"
	case UBLK_IO_OP_FLUSH:
		return "flush"
	case UBLK_IO_OP_DISCARD:
		return "discard"
	default:
		return "unknown"
	}
}

// Close closes the driver
func (d *Driver) Close() error {
	d.mu.Lock()
	defer d.mu.Unlock()

	ctx := context.Background()
	for id := range d.devices {
		d.RemoveDevice(ctx, id)
	}

	return nil
}

// Mock implementation note:
// This is a simplified implementation. A production ublk driver would need:
// 1. Proper syscall wrappers for ublk ioctls
// 2. Full io_uring integration with mmap'd ring buffers
// 3. Proper error handling and recovery
// 4. Performance optimizations (zero-copy, batch processing)
// 5. Support for all block operations (trim, write-same, etc.)
